from __future__ import with_statement
from __future__ import absolute_import

"""The core of the migration system"""

from ConfigParser import ConfigParser
from contextlib import closing
from optparse import OptionParser
import os
import re
import sys
import traceback
from unittest import TestCase

from sqlalchemy import create_engine

from dbmigrate.datasource import MigrationOptions
from dbmigrate.datasource.sqlite import DataSource as SQLiteDataSource
from dbmigrate.datasource.mysql import DataSource as MySQLDataSource
from dbmigrate.datasource.postgres import DataSource as PostgresDataSource
from dbmigrate.quarks import *
from dbmigrate.quarks import Atom, BaseAtom, Quark


__docformat__ = "restructuredtext"

__all__ = ['get_data_source', 'datasources', 'SucceedIf', 'SucceedUnless', 'Migration', 'Migrate', 'MigrationEngine', 'DBTestCase',
              'AdminConnection', 'MigrationConnection', 'AppConnection', 'RawConnection', 'CountIsZero',
              'WithCleanConnection', 'CreateDatabase', 'MigrationSucceeded', 'SucceedIf', 'MigrationWasApplied']


datasources = {SQLiteDataSource.scheme: SQLiteDataSource,
               MySQLDataSource.scheme: MySQLDataSource,
               PostgresDataSource.scheme: PostgresDataSource}

def get_data_source(scheme):
    return datasources[scheme]

class Migrate(object):

    """This is the main class that runs the migrations."""

    def __init__(self, args=None):
        """Set everything up, but don't run the migrations.

        args
          This defaults to ``sys.argv[1:]``.

        """
        self.setup_option_parser(args)
        self.eng = MigrationEngine(self.package, self.get_data_source(), self.options)

    def main(self):
        if self.options.print_revision:
            self.eng.print_revisions()
        else:
            self.eng.migrate_database()
            
    def get_data_source(self):
        ds = get_data_source(self.scheme)
        return ds.from_options(self.options, self.args)

    def setup_option_parser(self, args):
        """Parse command line arguments."""
        self.args = args
        usage = "usage: %prog [options] PACKAGE SCHEME [args]"
        self.parser = OptionParser(usage=usage)
        self.parser.add_option('-r', '--revision', type='int',
                          help='schema revision; defaults to most current')
        self.parser.add_option('-s', '--scheme',
                               help="the database URI scheme")
        self.parser.add_option('-p', '--print_revision', action="store_true",
                               default=False, 
                               help='print current revision and exit')
        self.parser.add_option("-v", "--verbose", action="store_const",
                               const=2, default=1)
        self.parser.add_option("--admin_user", default=None,
                               help="the database administrator")
        self.parser.add_option("--admin_pw", default=None,
                               help="the database administrator password")
        self.parser.add_option("--user", default=None,
                               help="the database user the application will use")
        self.parser.add_option("--pw", default=None,
                               help="the password the application user use")
        self.parser.add_option("--db", default=None,
                                help="the database to construct")
        self.parser.add_option("--host", "--server", default="localhost",
                               help="the database server")
        self.parser.add_option("--port", default=None,
                               help="the database server port")
#        self.parser.add_option("--socket", default=None,
#                               help="the unix domain socket for the database server")
        (self.options, self.args) = self.parser.parse_args(self.args)
        if len(self.args) < 2:
            self.parser.error("Expected at least one argument for PACKAGE and one for SCHEME")
        self.package =  self.args[0]
        self.scheme = self.args[1]
        self.args = self.args[2:]


class DBTestCase(TestCase):

    def setUp(self):
        super(DBTestCase, self).setUp()
        self._readIniFile()
        ds = get_data_source(self.scheme)
        self._datasource = ds.from_config_file(self._config)
        self._datasource.read_command(self)
        self._datasource.set_up_unit_test()
        self._options = self._optionsFromConfig()
        self._eng = MigrationEngine(self.package, self._datasource, self._options)
        try:
            self._options.revision = None
            self._eng.migrate_database()
        except:
            self._dropDatabase()
            raise
        self.connect_application(self._datasource.app_uri(), self._datasource.expansions)
        
    def tearDown(self):
        try:
            self.disconnect_application(self._datasource.app_uri(), self._datasource.expansions)
        except Exception:
            pass
        try:
            self._dropDatabase()
        finally:
            super(DBTestCase, self).tearDown()

    def _dropDatabase(self):
        if self._shouldKeepDatabase():
            self._datasource.print_unit_test_information()
        else:
            self._options.revision = 0
            self._eng.migrate_database()
            self._datasource.tear_down_unit_test()
           
    def _shouldKeepDatabase(self):
        if hasattr(self, 'keepDatabase') and self.keepDatabase:
            return True
        elif hasattr(self, 'keep_database') and self.keep_database:
            return True
        else:
            return False
        
    def _readIniFile(self):
        self._config = ConfigParser()
        self._configFile = self._findFile(os.getcwd(), 'setup.cfg')
        self._config.read(self._configFile)
        self._assertConfigHasSection()
        self._assertConfigHasPackage()
        self._assertConfigHasScheme()
        if hasattr(self, 'migrations'):
            self.package = self.migrations
        if not hasattr(self, 'package'):
            self.package = self._config.get('dbmigrate', 'package')
        if not hasattr(self, 'scheme'):
            self.scheme = self._config.get('dbmigrate', 'scheme')

    def _findFile(self, directory, name):
        fileHere = os.path.join(directory, name)
        if os.path.exists(fileHere):
            return fileHere
        else:
            parentDirectory = os.path.dirname(directory)
            if parentDirectory ==  directory:
                return None
            else:
                return self._findFile(parentDirectory, name)

    def _assertConfigHasSection(self):
        if not self._config.has_section('dbmigrate'):
            raise ValueError('The config file %s must define a dbmigrate section' % self._configFile)
 
    def _assertConfigHasPackage(self):
        if not self._config.has_option('dbmigrate', 'package'):
            raise ValueError('The config file %s must define the "package" option' % self._configFile)

    def _assertConfigHasScheme(self):
        if not self._config.has_option('dbmigrate', 'scheme'):
            raise ValueError('The config file %s must define the "scheme" option' % self._configFile)

    def _optionsFromConfig(self):
        return MigrationOptions(verbose=self._getVerbose(), print_revision=False, revision=None)

    def _getVerbose(self):
        return bool(self._getFromConfig('verbose', 'False'))
    
    def _getFromConfig(self, name, default):
        if self._config.has_option('dbmigrate', name):
            return self._config.get('dbmigrate', name)
        else:
            return default

    def connect_application(self, uri, expansions):
        """This function connects your application to the database.
        """
        pass
    
    def disconnect_application(self, uri, expansions):
        """This function disconnects your application from the database.
        """
        pass


class MigrationEngine(object):

    def __init__(self, package=None, datasource=None, options=None):
        self.package = package
        self.datasource = datasource
        self.options = options
        self.migrations = {}

    def run(self, params=None, packages=None, action=None):
        """Your application runs this to start the engine"""
        self.prepare(params, packages, action)
        if action is None:
            action = self.migrate_database
        action()
        
    def prepare(self, params=None, packages=None):
        if params is None:
            params = {}
        params['revision'] = params.get('revision', None)
        params['verbose'] = params.get('verbose', 1) 
        if packages is not None:
            self.package = packages
        if self.datasource is None:
            if 'scheme' not in params:
                raise KeyError('Scheme is a required parameter')
            ds = get_data_source(params['scheme'])
            self.datasource = ds.from_command(MigrationOptions(**params))
        if self.options is None:
            self.options = MigrationOptions(**params)
        
    def migrate_database(self):
        def migrate_action():
            self.print_overview()
            self.apply_downward_migrations()
            self.apply_upward_migrations()
        self.perform_action(migrate_action)

    def print_revisions(self):
        def print_revision():
            print "current revision: ",  self.current_revision()
            print "applied revisions: ", ', '.join([str(x) for x in self.applied_revisions])
            print "unapplied revisions: ", ', '.join([str(x.revision) for x in self.upward_migrations])
        self.perform_action(print_revision)

    def perform_action(self, action):
        """Run the migrations.

        All database activity starts from here.

        """
        with closing(self.datasource.connect()):
            self.find_migration_modules()
            self.find_desired_revision()
            self.find_applied_migrations()
            self.find_migrations_to_apply()
            action()

    def find_migration_modules(self):
        """Figure out what migrations exist."""
        self.migration_modules = self.get_migration_modules()
        if self.options.verbose:
            migration_files = [x.module.__file__ for x in self.migration_modules]
            print "The following migration files were found:", ', '.join(migration_files)

    def is_verbose(self):
        return self.options.verbose > 1
    
    def get_migration_modules(self):
        migrations = {}
        if not isinstance(self.package, list):
            self.package = [self.package]
        for package in self.package:
            self.load_migrations_from_files(migrations, package, self.module_files(package))
        return sorted(list(migrations.values()))            
    
    def module_files(self, package_name):
        module = __import__(package_name, [], [], [package_name])
        dirname = os.path.dirname(module.__file__)
        files = os.listdir(dirname)
        return [f for f in files if Migration.recognizesFile(f)]

    def load_migrations_from_files(self, migrations, package, files):
        for f in files:
            migration = Migration(package, f)
            self.assert_revision_is_unique(migration, migrations)
            migration.load()
            migrations[migration.revision] = migration
    
    def assert_revision_is_unique(self, migration, migrations):
        if migration.revision in migrations:
            msg = "migration revisions must be unique:  %s conflicts with %s"
            raise KeyError(msg % (migration, migrations[migration.revision]))
                           
    def find_desired_revision(self):
        """Find the target revision."""
        if self.options.revision is None:
            self.desired_revision = self.largest_migration_revision()
        else:
            self.desired_revision = self.options.revision

    def largest_migration_revision(self):
        indexes = sorted(self.migration_modules)
        return indexes[-1].revision

    def find_applied_migrations(self):
        """Figure out what revision we're currently at."""
        if self.datasource.has_schema_table():
            self.applied_revisions = self.datasource.get_applied_revisions()
        else:
            self.applied_revisions = []

    def find_migrations_to_apply(self):
        """Figure out the range of the migrations that need to be applied."""
        self.upward_migrations = sorted(self.choose_migrations(self.should_be_applied))
        self.downward_migrations = sorted(self.choose_migrations(self.should_be_unapplied), reverse=True)

    def should_be_applied(self, x):
        return x not in self.applied_revisions and x <= self.desired_revision

    def should_be_unapplied(self, x):
        return x in self.applied_revisions and x > self.desired_revision

    def choose_migrations(self, select):
        # choose migrations modules based on revision
        return [x for x in self.migration_modules if select(x.revision)]

    def print_overview(self):
        if self.is_verbose():
            subst = {'current': self.current_revision(), 'desired': self.desired_revision}
            if subst['desired'] == subst['current']:
                print "Updating current revision level %(current)s" % subst
            elif subst['desired'] > subst['current']:
                print "Upgrade from revision %(current)s to revision %(desired)s" % subst
            else:
                print "Downgrade from %(current)s to revision %(desired)s" % subst
                
    def print_overview_detailed(self):
        """If verbose, tell the user what's going on."""
        if self.is_verbose():
            print "Current revision:", self.current_revision()
            print "Desired revision:", self.desired_revision
            print "Migrations to be applied:", ', '.join([x.package_name() for x in self.upward_migrations])
            print "Migrations to be unapplied:", ', '.join([x.package_name() for x in self.downward_migrations])

    def current_revision(self):
        if self.applied_revisions:
            return self.applied_revisions[-1]
        else:
            return 0

    def apply_upward_migrations(self):
        for x in self.upward_migrations:
            self.apply_migration(self.apply_migration_up, "upward", x)

    def apply_downward_migrations(self):
        for x in self.downward_migrations:
            self.apply_migration(self.apply_migration_down, "downward", x)

    def apply_migration(self, application_function, direction, migration):
        """Apply the given migration list."""
        if self.is_verbose():
            print "Applying %s" % migration.package_name()
        application_function(migration.atoms(self.datasource.scheme), migration.revision, migration.name)

    def apply_migration_up(self, atoms, revision, name):
        """Apply the up steps in a migration, and the applied ones out if any fail"""
        undo_quarks =[]
        try:
            self.apply_up_atoms(atoms, revision, name, undo_quarks)
        except:
            print >> sys.stderr, "An exception occurred:"
            traceback.print_exc()
            print >> sys.stderr, "Trying to back out migration:", name
            self.undo_applied_atoms(undo_quarks)
            print >> sys.stderr, "Backed out migration:", name
            print >> sys.stderr, "Re-raising original exception."
            raise

    def apply_up_atoms(self, atoms, revision, name, undo_quarks):
        conn = self.datasource.migration_conn
        with conn.manage_transaction(revision):
            try:
                for x in atoms:
                    x.up(self.datasource, conn)
                    undo_quarks.append(x.down)
            except MigrationSucceeded:
                pass
            self.datasource.record_revision_applied(revision, name)
            
    def undo_applied_atoms(self, undo_quarks):
        conn = self.datasource.migration_conn
        with conn.manage_transaction():
            for down in reversed(undo_quarks):
                down(self.datasource, conn)

    def apply_migration_down(self, atoms, revision, name):
        """Apply the down steps of a migration"""
        with self.datasource.migration_conn.manage_transaction(revision):
            for x in reversed(atoms):
                x.down(self.datasource, self.datasource.migration_conn)
            self.datasource.record_revision_unapplied(revision, name)


class Migration(object):
    
    @classmethod
    def recognizesFile(cls, filename):
        return re.search('^\d+_.+\.py$', filename)
    
    def __init__(self, package, filename):
        self.package = package
        self.filename = filename
        parts = re.search('^((\d+)_.*)\.py', filename)
        self.name = parts.group(1)
        self.revision = int(parts.group(2))

    def load(self):
        self.module = __import__(self.package_name(), [], [], [self.name])
        self.migrations = self.module.migrations 

    def package_name(self):
        return '%s.%s' % (self.package, self.name)
    
    def __cmp__(self, x):
        return long(self.revision).__cmp__(long(x))
    
    def __hash__(self):
        return self.revision

    def __long__(self):
        return long(self.revision)

    def atoms(self, scheme):
        return [self.atom_factory(x) for x in self.raw_atoms(scheme)]

    def raw_atoms(self, scheme):
        if isinstance(self.migrations, dict):
            return self.raw_atoms_from_dict(scheme)
        else:
            return self.migrations
        
    def raw_atoms_from_dict(self, scheme):
        if scheme in self.migrations:
            return self.migrations[scheme]
        elif '_' in self.migrations:
            return self.migrations['_']
        else:
            return []
        
    def atom_factory(self, x):
        if isinstance(x, BaseAtom):
            return x
        else:
            return Atom(x)


def main():
    Migrate().main()
